from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import InMemorySaver 


class State(TypedDict):
    input: str


def login_email(state):
    print("---Step 1  登录邮件---")
    pass


def edit_email(state):
    print("---Step 2  编辑邮件---")
    pass


def send_email(state):
    print("---Step 3  发送邮件---")
    pass


builder = StateGraph(State)
builder.add_node("login_email", login_email)
builder.add_node("edit_email", edit_email)
builder.add_node("send_email", send_email)
builder.add_edge(START, "login_email")
builder.add_edge("login_email", "edit_email")
builder.add_edge("edit_email", "send_email")
builder.add_edge("send_email", END)

# Set up a checkpointer 
checkpointer = InMemorySaver() # (1)!

graph = builder.compile(
    checkpointer=checkpointer, # (2)!
    # interrupt_after=["step_1"] # 执行step1后，中断!
    interrupt_before=["send_email"], # 执行step3前，中断!
)


# Input
initial_input = {"input": "hello world"}

# Thread
thread = {"configurable": {"thread_id": "1"
""
""}}

# Run the graph until the first interruption
for event in graph.stream(initial_input, thread, stream_mode="values"):
    print(event)

# This will run until the breakpoint
# You can get the state of the graph at this point
print(graph.get_state(thread))



ack_input = "yes"

# 可以输入 `None` 恢复流程 
for event in graph.stream(ack_input, thread, stream_mode="values"):
    print(event)